import numpy as np

def im2col(img, ksize, stride=1):
    N, H, W, C = img.shape
    # 等于卷积输出图像的高，也等于纵向卷积的次数
    out_h = (H - ksize) // stride + 1
    # 等于卷积输出图像的宽，也等于横向卷积的次数
    out_w = (W - ksize) // stride + 1
    # 每次卷积要用到的图像排成一行，每行的个数为卷积核的尺寸×通道数;out_h * out_w表示对于一幅图像卷积需要的次数；
    col = np.empty((N * out_h * out_w, ksize * ksize * C))
    outsize = out_w * out_h
    # 按照卷积的过程展开
    for y in range(out_h):
        y_min = y * stride
        y_max = y_min + ksize
        y_start = y * out_w
        for x in range(out_w):
            x_min = x * stride
            x_max = x_min + ksize
            col[y_start + x::outsize, :] = img[:, y_min:y_max, x_min:x_max, :].reshape(N, -1)
    print("after apply img3col:\n",col)
    return col


def conv(X, kernel, stride=1, padding='same'):
    FN, ksize, ksize, C = kernel.shape
    if padding == 'same':
        p = ksize // 2
        X = np.pad(X, ((0, 0), (p, p), (p, p), (0, 0)), 'constant')
    print("after padding(x): ",X)
    N, H, W, C = X.shape
    col = im2col(X, ksize, stride)
    z = np.dot(col, kernel.reshape(kernel.shape[0], -1).transpose())
    z = z.reshape(N, z.shape[0] // N, -1)
    out_h = (H - ksize) // stride + 1
    return z.reshape(N, out_h, -1, FN)

if __name__ == "__main__":
    testArry = np.array([[[[1],[2],[3],[4]],
                          [[5],[6],[7],[8]],
                          [[9],[10],[11],[12]],
                          [[13],[14],[15],[16]]]])
    kernel = np.array([[[[1],[1],[1]],
                   [[1],[1],[1]],
                   [[1],[1],[1]]]])
    # print(W.reshape(W.shape[0], -1).transpose())
    # N, H, W, C = testArry.shape
    # print("NHWC: ",N, H, W, C)
    # col = im2col(testArry,3)
    # print("col: ",col)

    convRes = conv(testArry,kernel)
    print("conv result: \n",convRes)